import asyncio
import argparse
import subprocess
import sys
from pathlib import Path
from rosbags.rosbag1 import Reader as Rosbag1Reader, Writer as Rosbag1Writer
from rosbags.rosbag1.reader import ReaderError
from rosbags.typesys import get_typestore, Stores, get_types_from_msg
from rosbags.serde import serialize_ros1
from typing import Optional, List, Dict, Any
from .customchange import  process_custom_conversions_for_bag

try:
    import numpy as np
    NUMPY_AVAILABLE = True
except ImportError:
    NUMPY_AVAILABLE = False

# 尝试导入项目专用日志记录器和配置管理器
try:
    from gglobal import logger,config

except ImportError:
    import logging
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)
    handler = logging.StreamHandler()
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    from gglobal.param_manager import ConfigManager
    config = ConfigManager()

class RosBagConverter:
    def __init__(self, bag_path: Path, output_bag_path: Optional[Path] = None, target_topics: Optional[List[str]] = None, custom_msg_path: Optional[Path] = None):
        """初始化ROS包转换器

        参数:
            bag_path: 输入ROS包文件路径
            output_bag_path: 可选，输出bag文件路径
            target_topics: 可选，要保存的主题列表
            custom_msg_path: 可选，自定义消息类型定义文件所在的根目录路径
        """
        self.bag_path = bag_path
        
        if output_bag_path:
            self.output_bag_path = Path(output_bag_path)
        else:
            # 默认输出路径：在原文件同目录下创建 <原文件名>_converted.bag
            self.output_bag_path = self.bag_path.parent / f"{self.bag_path.stem}_converted.bag"
            
        self.target_topics = target_topics
        self.custom_msg_path = custom_msg_path
        self.cancel_requested = False
        self.current_message_count = 0  # 当前已转换的消息条数

        # 初始化类型存储
        try:
            self.typestore = get_typestore(Stores.ROS1_NOETIC)
        except Exception as e:
            logger.warning(f"初始化默认ROS1类型存储失败：{e}。将尝试重新创建。")
            self.typestore = get_typestore(Stores.ROS1_NOETIC)
        
        if self.custom_msg_path:
            self._register_custom_types(self.custom_msg_path)
        else:
            # 提供一个默认的自定义消息路径，如果用户没有指定
            default_msg_dir = Path(__file__).parent / 'msg'
            if default_msg_dir.exists():
                logger.info(f"未指定自定义消息路径，尝试使用默认路径: {default_msg_dir}")
                self._register_custom_types(default_msg_dir)
            else:
                logger.info("未指定自定义消息路径，且默认路径不存在，将仅使用内置类型。")

    def _register_custom_types(self, msg_root_path: Path):
        """注册所有自定义消息类型。"""
        # 定义消息类型的优先级（只针对自定义类型）
        priority = {
            "Point32.msg": 0,  # geometry_msgs/Point32 需要先注册
            "Mbr.msg": 1,      # mooe_msgs/Mbr 依赖 Point32
            "Obstacle.msg": 2, # mooe_msgs/Obstacle 依赖 Mbr 和 Point32
            "ObstacleArray.msg": 3  # mooe_msgs/ObstacleArray 依赖 Obstacle
        }

        # 收集所有.msg文件
        msg_files = list(msg_root_path.rglob('*.msg'))

        # 根据优先级排序
        msg_files.sort(key=lambda p: priority.get(p.name, 100))
        logger.info(f"消息文件注册顺序: {[f.name for f in msg_files]}")

        for msg_file in msg_files:
            pkg = msg_file.parent.name
            name = msg_file.stem
            typename = f'{pkg}/{name}'

            try:
                msg_text = msg_file.read_text(encoding='utf-8')
                
                # 检查bagversion参数，如果是old版本且是Obstacle.msg，则移除category字段
                if name == 'Obstacle' and config is not None:
                    try:
                        bagversion = config.get_param('bagversion', 'new')
                        if bagversion == 'old':
                            # 移除category字段行
                            lines = msg_text.split('\n')
                            filtered_lines = [line for line in lines if not line.strip().startswith('uint32 category')]
                            msg_text = '\n'.join(filtered_lines)
                            logger.info(f"检测到bagversion为old，已从{typename}中移除category字段")
                    except Exception as e:
                        logger.warning(f"获取bagversion参数失败: {e}，使用默认处理")
                
                # 检查类型是否已存在
                if typename in self.typestore.fielddefs:
                    logger.warning(f"类型 {typename} 已存在，跳过注册。")
                    continue

                # 从.msg文件内容中获取类型定义
                custom_types = get_types_from_msg(msg_text, f'{pkg}/{name}')
                
                # 注册类型
                self.typestore.register(custom_types)
                logger.info(f'成功注册自定义消息类型: {typename}')
            except Exception as e:
                logger.error(f'注册消息类型 {typename} 失败: {e}')
        
        logger.debug(f"当前已注册的所有类型: {list(self.typestore.fielddefs.keys())}")
    
    def _msg_to_dict(self, msg) -> Dict[str, Any]:
        """自适应地将所有 ROS 数据类型转换为字典"""
        # 检查是否为基本类型或非 ROS 对象
        if not hasattr(msg, '__msgtype__'):
            if isinstance(msg, (str, int, float, bool)):
                return msg
            elif isinstance(msg, (bytes, bytearray)):
                return msg.hex()  # 转换为十六进制字符串
            elif NUMPY_AVAILABLE and isinstance(msg, np.ndarray):
                return msg.tolist()  # 转换为列表
            elif msg is None:
                return None
            else:
                logger.warning(f"未知类型 {type(msg)}，值：{msg}，转为字符串")
                return str(msg)

        # 处理 ROS 消息对象
        result = {}
        for attr in dir(msg):
            if (attr.startswith('_') or 
                attr in ['__msgtype__', '_md5sum', '_type', 'get_fields_and_field_types', 
                         'SLOT_TYPES', 'get_field_names']):
                continue
            
            try:
                value = getattr(msg, attr)
            except AttributeError:
                logger.warning(f"无法获取属性 {attr}，跳过")
                continue

            # 处理列表或数组
            if isinstance(value, (list, tuple)):
                result[attr] = [self._msg_to_dict(item) for item in value]
            # 处理 NumPy 数组（多维数组）
            elif NUMPY_AVAILABLE and isinstance(value, np.ndarray):
                result[attr] = value.tolist()
            # 处理嵌套 ROS 消息对象
            elif hasattr(value, '__msgtype__'):
                result[attr] = self._msg_to_dict(value)
            # 处理时间类型（Time 或 Duration）
            elif hasattr(value, 'sec') and hasattr(value, 'nanosec'):
                result[attr] = {
                    "sec": value.sec,
                    "nanosec": value.nanosec
                }
            # 基本类型
            else:
                result[attr] = value
        
        return result
    
    async def _process_custom_conversions(self, topic_name, msg_dict, timestamp, writer, generated_connections):
        """处理自定义话题转换"""
        try:
            # 调用customchange模块中的转换函数
            from .customchange import process_custom_conversions_for_bag
            await process_custom_conversions_for_bag(
                topic_name, msg_dict, timestamp, writer, generated_connections, 
                self.typestore, self._dict_to_msg
            )
        except Exception as e:
            self.logger.error(f"处理自定义转换时出错: {e}", exc_info=True)
    

    def _dict_to_msg(self, data: dict, msgdef) -> Any:
        """将字典转换为ROS消息对象"""
        try:
            constructor_args = self._prepare_default_args(msgdef)
            constructor_args.update(self._process_field_values(data, msgdef))
            return self._create_message_object(msgdef, constructor_args)
        except Exception as e:
            logger.error(f"字典转换为消息对象失败: {e}")
            logger.error(f"失败的数据类型: {type(data)}, 消息类型: {msgdef.cls.__name__}")
            logger.debug(f"消息定义: {msgdef}")
            raise

    def _prepare_default_args(self, msgdef) -> Dict[str, Any]:
        """为所有字段准备默认值"""
        constructor_args = {}
        
        for field_tuple in msgdef.fields:
            field_name = field_tuple[0]
            field_info = field_tuple[1]
            
            # 根据字段类型设置默认值
            if len(field_info) > 1:
                field_type = field_info[1]
                constructor_args[field_name] = self._get_default_value_for_type(field_type)
            else:
                constructor_args[field_name] = None
                
        return constructor_args
    
    def _get_default_value_for_type(self, field_type) -> Any:
        """根据字段类型获取默认值"""
        if isinstance(field_type, tuple):
            # 基本类型 (type_name, size)
            type_name = field_type[0]
            if type_name == 'bool':
                return False
            elif type_name in ['int8', 'int16', 'int32', 'int64', 'uint8', 'uint16', 'uint32', 'uint64']:
                return 0
            elif type_name in ['float32', 'float64']:
                return 0.0
            elif type_name == 'string':
                return ''
            else:
                return None
        elif isinstance(field_type, str):
            # 嵌套消息类型，创建默认消息对象
            try:
                nested_msgdef = self.typestore.get_msgdef(field_type)
                return self._dict_to_msg({}, nested_msgdef)
            except Exception:
                return None
        else:
            return None

    def _process_field_values(self, data: dict, msgdef) -> Dict[str, Any]:
        """处理实际字段值"""
        processed_args = {}
        
        for field_name, field_value in data.items():
            field_info = self._get_field_info(msgdef, field_name)
            
            if field_info:
                field_type = field_info[1] if len(field_info) > 1 else None
                processed_args[field_name] = self._convert_field_value(field_value, field_type)
                
        return processed_args
    
    def _get_field_info(self, msgdef, field_name: str):
        """获取字段类型信息"""
        for field_tuple in msgdef.fields:
            if field_tuple[0] == field_name:
                return field_tuple[1]
        return None
    
    def _convert_field_value(self, field_value: Any, field_type) -> Any:
        """转换字段值"""
        if isinstance(field_value, dict) and field_type:
            # 嵌套消息对象
            return self._convert_nested_message(field_value, field_type)
        elif isinstance(field_value, list) and field_type:
            # 处理数组字段
            return self._convert_array_field(field_value, field_type)
        else:
            # 基本类型字段
            return field_value
    
    def _convert_nested_message(self, field_value: dict, field_type) -> Any:
        """转换嵌套消息对象"""
        try:
            nested_msgdef = self.typestore.get_msgdef(field_type)
            return self._dict_to_msg(field_value, nested_msgdef)
        except Exception:
            # 如果无法获取嵌套类型定义，直接设置值
            return field_value
    
    def _convert_array_field(self, field_value: list, field_type) -> list:
        """转换数组字段"""
        if isinstance(field_type, tuple) and len(field_type) > 0:
            # 序列类型，获取元素类型
            element_type_info = field_type[0]
            if isinstance(element_type_info, tuple) and len(element_type_info) > 1:
                element_type = element_type_info[1]
                return self._convert_array_elements(field_value, element_type)
            else:
                return field_value
        else:
            return field_value
    
    def _convert_array_elements(self, field_value: list, element_type: str) -> list:
        """转换数组元素"""
        converted_list = []
        for item in field_value:
            if isinstance(item, dict):
                try:
                    item_msgdef = self.typestore.get_msgdef(element_type)
                    converted_list.append(self._dict_to_msg(item, item_msgdef))
                except Exception:
                    converted_list.append(item)
            else:
                converted_list.append(item)
        return converted_list
    
    def _create_message_object(self, msgdef, constructor_args: Dict[str, Any]) -> Any:
        """创建消息对象"""
        return msgdef.cls(**constructor_args)
    
    def _validate_and_reindex_bag(self):
        """验证bag文件索引，如果需要则重新索引"""
        try:
            # 尝试打开bag文件检查索引
            with Rosbag1Reader(self.bag_path) as reader:
                # 如果能成功打开，说明索引正常
                logger.debug(f"Bag文件索引正常: {self.bag_path}")
                return
        except ReaderError as e:
            if "not indexed" in str(e).lower():
                logger.warning(f"检测到bag文件未索引，正在重新索引: {self.bag_path}")
                self._reindex_bag()
            else:
                # 其他类型的错误，重新抛出
                raise
        except Exception as e:
            logger.error(f"验证bag文件时发生未知错误: {e}")
            raise
    
    def _reindex_bag(self):
        """重新索引bag文件"""
        try:
            # 使用rosbag reindex命令重新索引
            cmd = ["rosbag", "reindex", str(self.bag_path)]
            logger.info(f"执行重索引命令: {' '.join(cmd)}")
            
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=300  # 5分钟超时
            )
            
            if result.returncode == 0:
                logger.info(f"Bag文件重索引成功: {self.bag_path}")
            else:
                error_msg = f"重索引失败，返回码: {result.returncode}\n标准输出: {result.stdout}\n错误输出: {result.stderr}"
                logger.error(error_msg)
                raise RuntimeError(f"重索引bag文件失败: {error_msg}")
                
        except subprocess.TimeoutExpired:
            error_msg = f"重索引操作超时: {self.bag_path}"
            logger.error(error_msg)
            raise RuntimeError(error_msg)
        except FileNotFoundError:
            error_msg = "未找到rosbag命令，请确保ROS环境已正确安装和配置"
            logger.error(error_msg)
            raise RuntimeError(error_msg)
        except Exception as e:
            error_msg = f"重索引过程中发生错误: {e}"
            logger.error(error_msg)
            raise RuntimeError(error_msg)


    async def process_bag(self):
        """异步读取ROS包文件并将数据写入新的bag文件。"""
        if not self.bag_path.exists():
            logger.error(f"未找到包文件：{self.bag_path}")
            return

        # 确保输出目录存在
        self.output_bag_path.parent.mkdir(parents=True, exist_ok=True)
        
        # 如果输出文件已存在，删除它
        if self.output_bag_path.exists():
            self.output_bag_path.unlink()
            logger.info(f"已删除现有输出文件：{self.output_bag_path}")

        logger.info(f"开始处理bag文件: {self.bag_path.name}")
        
        # 检查并重索引bag文件（如果需要）
        self._validate_and_reindex_bag()

        try:
            with Rosbag1Reader(self.bag_path) as reader:
                logger.info(f"正在处理包文件：{self.bag_path}")
                logger.info(f"包文件时长：{reader.duration * 1e-9:.2f}秒，开始时间：{reader.start_time}，结束时间：{reader.end_time}，消息数量：{reader.message_count}")
                
                # 收集要处理的连接信息
                connections_to_process = []
                for conn_obj in reader.connections:
                    topic_name = conn_obj.topic
                    # 如果设置了目标主题，则仅处理这些主题
                    if self.target_topics and topic_name not in self.target_topics:
                        logger.info(f"跳过非目标主题: {topic_name}")
                        continue
                    connections_to_process.append(conn_obj)
                
                logger.info(f"找到 {len(connections_to_process)} 个主题需要处理")
                
                # 创建输出bag文件
                with Rosbag1Writer(self.output_bag_path) as writer:
                    # 添加连接信息到输出bag
                    connection_map = {}
                    # 用于跟踪已添加的连接，避免重复添加相同的topic+msgtype组合
                    added_connections = {}
                    
                    for conn_obj in connections_to_process:
                        # 创建唯一标识符
                        connection_key = (conn_obj.topic, conn_obj.msgtype)
                        
                        # 检查是否已经添加过相同的连接
                        if connection_key in added_connections:
                            # 如果已存在，复用已有的输出连接
                            output_conn = added_connections[connection_key]
                            logger.debug(f"复用现有连接: {conn_obj.topic} ({conn_obj.msgtype})")
                        else:
                            # 在输出bag中添加新连接
                            output_conn = writer.add_connection(
                                topic=conn_obj.topic,
                                msgtype=conn_obj.msgtype,
                                typestore=self.typestore
                            )
                            added_connections[connection_key] = output_conn
                            logger.debug(f"添加连接: {conn_obj.topic} ({conn_obj.msgtype})")
                        
                        connection_map[conn_obj.id] = output_conn
                    
                    # 处理消息
                    count = 0
                    
                    # 用于存储新生成的连接
                    generated_connections = {}
                    
                    for connection, timestamp, rawdata in reader.messages():
                        # 检查取消请求
                        if self.cancel_requested:
                            logger.info("转换已被用户取消")
                            logger.info("转换已取消")
                            return
                        
                        topic_name = connection.topic
                        # 如果设置了目标主题，则仅处理这些主题的消息
                        if self.target_topics and topic_name not in self.target_topics:
                            continue
                        
                        # 检查连接是否在我们的映射中
                        if connection.id not in connection_map:
                            logger.warning(f"跳过未映射的连接ID: {connection.id} (主题: {topic_name})")
                            continue
                        
                        # 获取对应的输出连接
                        output_connection = connection_map[connection.id]
                        
                        # 写入原始消息
                        writer.write(output_connection, timestamp, rawdata)
                        count += 1
                        # 实时更新当前消息计数
                        self.current_message_count = count
                        
                        # 处理自定义话题转换
                        try:
                            # 反序列化消息以进行处理
                            msg = self.typestore.deserialize_ros1(rawdata, connection.msgtype)
                            msg_dict = self._msg_to_dict(msg)
                            
                            # # 处理自定义转换
                            await self._process_custom_conversions(topic_name, msg_dict, timestamp, writer, generated_connections)
                            
                        except KeyboardInterrupt:
                            raise
                        except Exception as e:
                            # 记录错误但不中断处理
                            logger.debug(f"处理话题 {topic_name} 的自定义转换时出错: {e}")
                        
                        # 更新当前消息计数
                        self.current_message_count = count
                
                if not self.cancel_requested:
                    logger.info(f"包文件处理完成。总共处理消息数：{count}")
                    logger.info(f"输出文件保存至：{self.output_bag_path}")
                else:
                    logger.info(f"转换已取消。已处理消息数：{count}")

        except Exception as e:
            error_msg = f"处理包文件{self.bag_path}时出错：{e}"
            logger.error(error_msg, exc_info=True)
            raise
    
    def cancel_conversion(self):
        """请求取消转换过程"""
        self.cancel_requested = True
        logger.info("已请求取消转换")

async def main_async():
    parser = argparse.ArgumentParser(description="将ROS1包文件转换为新的bag文件，可选择特定主题。")
    parser.add_argument("input_bag", type=str, nargs='?', 
                       default="F:/桌面/MooeMars/code/MooeTool/db/20250207135800.bag",
                       help="输入的ROS1 bag文件路径（默认：F:/桌面/MooeMars/code/MooeTool/db/20250207135800.bag）")
    parser.add_argument("--output_bag", type=str, help="可选：输出bag文件的路径。默认为当前路径下的bagchange.bag")
    parser.add_argument("--topics", type=str, nargs='*', help="可选：要处理的主题列表。如果不指定，则使用默认话题列表")
    parser.add_argument("--custom_msg_path", type=str, help="可选：自定义消息类型定义文件所在的根目录路径")
    args = parser.parse_args()

    bag_path = Path(args.input_bag)
    
    # 设置默认输出路径为当前路径下的bagchange.bag
    if args.output_bag:
        output_bag_path = Path(args.output_bag)
    else:
        output_bag_path = Path("bagchange.bag")
    
    # 设置默认目标话题列表
    if args.topics:
        target_topics = args.topics
    else:
        # 使用默认话题列表
        target_topics = [
            "/server_time_now",
            "/baselink_tf",
            "/lidar_perception/obstacle_array",
            "/scan_down_left",
            "/scan_down_right",
            "/tf",
        ]
    
    custom_msg_path = Path(args.custom_msg_path) if args.custom_msg_path else None

    if not bag_path.exists():
        logger.error(f"输入bag文件不存在: {bag_path}")
        return

    logger.info(f"输入文件: {bag_path}")
    logger.info(f"输出文件: {output_bag_path or bag_path.parent / f'{bag_path.stem}_converted.bag'}")
    if target_topics:
        logger.info(f"目标主题: {target_topics}")
    else:
        logger.info("将处理所有主题")

    converter = RosBagConverter(
        bag_path=bag_path,
        output_bag_path=output_bag_path,
        target_topics=target_topics,
        custom_msg_path=custom_msg_path
    )
    await converter.process_bag()

if __name__ == "__main__":
    asyncio.run(main_async())